import os

import gevent
from gevent.threadpool import ThreadPool
from gevent.fileobject import FileObjectThread

from app.utils.ws_data import *
from ws_state import WSState
from app_config import APP_USER_CONFIG
from utils.file_md5 import file_md5

from .fod.file_sync_fod import Config as FileSyncFODConfig, FileSyncFOD

from app_instance import app

from flask_sockets import Sockets
sockets = Sockets(app)

loop = gevent.get_hub().loop

def check_recv_data(data):
    if data is None:
        app.logger.error('data is None')
        return False
    if not isinstance(data, bytearray):
        app.logger.error('data type is not bytesarray')
        return False
    return True

class FileObjetReference:
    def __init__(self, file_path, mode, thread_pool):
        self.file_path = file_path
        self.mode = mode
        self.thread_pool = thread_pool
        self.fobject_ = FileObjectThread(open(self.file_path, self.mode), threadpool=self.thread_pool)
        app.logger.debug('file %s opened' % self.file_path)
        self.ref_count = 0

    def ref(self):
        self.ref_count += 1
        if self.fobject_.closed:
            self.fobject_ = FileObjectThread(open(self.file_path, self.mode), threadpool=self.thread_pool)
            app.logger.debug('file %s opened' % self.file_path)

    def unref(self):
        self.ref_count -= 1
        if self.ref_count <= 0:
            self.close()
            app.logger.debug('file %s closed' % self.file_path)

    def closed(self):
        return self.ref_count <= 0

    def close(self):
        self.ref_count = 0
        self.fobject_.close()

    def fobject(self):
        return self.fobject_

class CachedFileObject:
    def __init__(self):
        self.cached_fobjects = {}

    def ref(self, file_name, file_path, mode, thread_pool):
        fobj_ref = None
        if file_name in self.cached_fobjects:
            fobj_ref = self.cached_fobjects[file_name]
        else:
            fobj_ref = FileObjetReference(file_path, mode, thread_pool)
            self.cached_fobjects[file_name] = fobj_ref
        assert fobj_ref is not None
        fobj_ref.ref()
        return fobj_ref.fobject()

    def unref(self, file_name):
        if file_name in self.cached_fobjects:
            fobj_ref = self.cached_fobjects[file_name]
            fobj_ref.unref()
            if fobj_ref.closed():
                self.cached_fobjects.__delitem__(file_name)

    def close(self, file_name):
        if file_name in self.cached_fobjects:
            fobj_ref = self.cached_fobjects[file_name]
            fobj_ref.close()
            self.cached_fobjects.__delitem__(file_name)

    def get(self, file_name):
        if file_name in self.cached_fobjects:
            return self.cached_fobjects[file_name].fobject()
        return None

class WSServer:
    '''websocket server'''
    def __init__(self,
                 client_id = APP_USER_CONFIG['clientId'],
                 data_dir = APP_USER_CONFIG['dataDir'],
                 inbox_dir = APP_USER_CONFIG['inbox'],
                 outbox_dir = APP_USER_CONFIG['outbox'],
                 chunk_size = APP_USER_CONFIG['chunkSize'],
                 task_timeout_secs = APP_USER_CONFIG['taskTimeoutSeconds']):
        enable_fod = False
        if 'enableFOD' in APP_USER_CONFIG:
            enable_fod = APP_USER_CONFIG['enableFOD']

        file_sync_fod_cfg = FileSyncFODConfig()
        if 'partSize' in APP_USER_CONFIG:
            file_sync_fod_cfg.part_size = APP_USER_CONFIG['partSize']
        if 'queueSize' in APP_USER_CONFIG:
            file_sync_fod_cfg.queue_size = APP_USER_CONFIG['queueSize']
        if 'costPerPart' in APP_USER_CONFIG:
            file_sync_fod_cfg.cost_per_part = APP_USER_CONFIG['costPerPart']
        if 'minCostTime' in APP_USER_CONFIG:
            file_sync_fod_cfg.min_cost_time = APP_USER_CONFIG['minCostTime']
        if 'maxCostTime' in APP_USER_CONFIG:
            file_sync_fod_cfg.max_cost_time = APP_USER_CONFIG['maxCostTime']
        if 'tmpPartInbox' in APP_USER_CONFIG:
            file_sync_fod_cfg.tmp_part_inbox = APP_USER_CONFIG['tmpPartInbox']
        file_sync_fod_cfg.part_inbox = inbox_dir
        file_sync_fod_cfg.part_outbox = outbox_dir
        file_sync_fod_cfg.data_dir = data_dir
        file_sync_fod_cfg.inbox = os.path.join(data_dir, '_inbox')
        file_sync_fod_cfg.outbox = os.path.join(data_dir, '_outbox')
        if enable_fod:
            self.file_sync_fod = FileSyncFOD(cfg=file_sync_fod_cfg, db_engine=app.db.engine)

        self.client_id = client_id
        if self.client_id is None:
            self.client_id = APP_USER_CONFIG['clientId']

        self.cached_fobjects = CachedFileObject()

        self.chunk_size = chunk_size

        self.thread_pool = ThreadPool(4)

        self.state = WSState(
            self.client_id,
            loop,
            data_dir=data_dir,
            inbox_dir=file_sync_fod_cfg.inbox if enable_fod else inbox_dir,
            outbox_dir=file_sync_fod_cfg.outbox if enable_fod else outbox_dir,
            logger=app.logger,
            task_timeout_secs=task_timeout_secs,
            is_server=True,
            chunk_size=chunk_size)

ws_server = WSServer()

@sockets.route('/websocket/stream/up')
def websocket_stream_up(ws):
    try:
        while not ws.closed:
            data = ws.receive()

            if ws.closed or not check_recv_data(data):
                continue

            opcode, header, data = WSDataPack.unpack(bytes(data))

            if opcode == WSOPCode.PUSH_FILE.value or opcode == WSOPCode.PUSH_CHUNKS.value:
                if 'file_name' not in header:
                    app.logger.error('file name not in header')
                    continue

                if 'size' not in header:
                    app.logger.error('size not in header')
                    continue

                if 'client_id' not in header:
                    app.logger.error('client id not in header')
                    continue

                client_id = header['client_id']
                file_name = header['file_name']
                file_size = header['size']
                checksum = header['checksum']

                if ws_server.state.is_finished_task(file_name):
                    app.logger.debug('adding finish task, file name: %s' % file_name)
                    ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': WSOPCode.TASK_FINISHED.value}, None))
                    continue

                if ws_server.state.is_push_task_exists(file_name) or \
                    (opcode == WSOPCode.PUSH_FILE.value and ws_server.state.is_unfinished_task(file_name)):
                        ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': WSOPCode.TASK_EXIST.value}, None))
                        continue

                file_path = ws_server.state.get_tmp_file_path(file_name)

                r = ws_server.state.add_unfinished_state({
                    'file_name': file_name,
                    'clients_id': [client_id],
                    'size': file_size,
                    'checksum': checksum
                    })

                ws_server.cached_fobjects.ref(file_name, file_path, 'wb+' if r else 'rb+',
                                                        ws_server.thread_pool)

                ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': opcode}, None))

                if file_size == 0:
                    ws_server.cached_fobjects.unref(file_name)

                app.logger.debug('ack send, file name: %s' % file_name)
            elif opcode == WSOPCode.PUSH_CHUNK.value:
                if 'file_name' not in header:
                    app.logger.error('file name not in header')
                    continue

                if 'offset' not in header:
                    app.logger.error('offset not in header')
                    continue

                file_name = header['file_name']
                offset = header['offset']
                chunk = data

                fobject = ws_server.cached_fobjects.get(file_name)
                if fobject:
                    fobject.seek(offset)
                    fobject.write(chunk)
                    fobject.flush()

                    ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': WSOPCode.PUSH_CHUNK.value}))

                    unfinished_task = ws_server.state.get_unfinished_state(file_name)

                    if ws_server.state.update_unfinished_state(file_name, (offset, len(chunk))):
                        ws_server.cached_fobjects.close(file_name)

                        assert unfinished_task is not None

                        file_name = unfinished_task['file_name']
                        tmp_file_path = ws_server.state.get_tmp_file_path(file_name)

                        try:
                            verify_succeed = True
                            checksum = unfinished_task['checksum']
                            app.logger.debug('verify file %s\'s checksum' % tmp_file_path)
                            md5 = file_md5(tmp_file_path)
                            if md5 != checksum:
                                app.logger.error('file %s\'s checksum not match' % tmp_file_path)
                                verify_succeed = False

                            if verify_succeed:
                                app.logger.debug('adding finish task, file name: %s' % file_name)
                                ws_server.state.add_finish_task_record(unfinished_task)
                                ws_server.state.add_transient_task(
                                    {'opcode': WSOPCode.FINISH_TASK.value, 'file_name': file_name,
                                     'clients_id': unfinished_task['clients_id']})

                            ws_server.state.finish_unfinished_state(file_name)

                        except Exception as ex:
                            app.logger.error('exception %s happend while check file checksum %s' % (ex, file_path))

            gevent.sleep(0.1)
    except Exception as ex:
        app.logger.error(ex)

@sockets.route('/websocket/stream/down')
def websocket_stream_down(ws):
    try:
        while not ws.closed:
            data = ws.receive()

            if ws.closed or not check_recv_data(data):
                continue

            opcode, header, data = WSDataPack.unpack(bytes(data))

            if 'client_id' not in header:
                app.logger.error('client id not in header')
                continue

            if 'file_name' not in header:
                app.logger.error('file_name not in header')
                continue

            file_name = header['file_name']
            client_id = header['client_id']

            send_ack = False

            file_path = ws_server.state.get_outbox_file_path(file_name)
            if os.path.exists(file_path):
                if opcode == WSOPCode.PULL_FILE.value:
                    app.logger.debug('pull file %s from client %s' % (file_name, client_id))
                    try:
                        fobject = ws_server.cached_fobjects.ref(file_name, file_path, 'rb',
                                                      ws_server.thread_pool)
                        ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': opcode}))
                        send_ack = True

                        offset = 0
                        while True:
                            fobject.seek(offset)
                            chunk = fobject.read(ws_server.chunk_size)
                            if chunk is not None and len(chunk) > 0:
                                ws.send(WSDataPack.pack(WSOPCode.PUSH_CHUNK.value, {'offset': offset}, chunk))

                                data = ws.receive()
                                if ws.closed or not check_recv_data(data):
                                    break
                                opcode, header, data = WSDataPack.unpack(bytes(data))
                                if opcode != WSOPCode.ACK.value or \
                                    'ack' not in header and \
                                    header['ack'] != WSOPCode.PUSH_CHUNK.value:
                                    break

                                offset += len(chunk)
                                ws_server.state.update_unack_task_modify_t(file_name)
                            else:
                                break
                            gevent.sleep(0.1)
                        app.logger.debug('send file %s to client %s finish' % (file_name, client_id))
                    except Exception as ex:
                        app.logger.error(ex)
                    finally:
                        ws_server.cached_fobjects.unref(file_name)

                elif opcode == WSOPCode.PULL_CHUNKS.value:
                    app.logger.debug('pull file %s chunks from %s' % (file_name, client_id))

                    if 'chunks' in header:
                        chunks = header['chunks']
                        try:
                            fobject = ws_server.cached_fobjects.ref(file_name, file_path, 'rb',
                                                                    ws_server.thread_pool)
                            ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': opcode}))
                            send_ack = True

                            for chunk_item in chunks:
                                offset = chunk_item[0]
                                dst_size = offset + chunk_item[1]
                                while offset < dst_size:
                                    fobject.seek(offset)
                                    chunk_data = fobject.read(ws_server.chunk_size)
                                    if chunk_data is not None and len(chunk_data) > 0:
                                        ws.send(WSDataPack.pack(WSOPCode.PUSH_CHUNK.value, {'offset': offset}, chunk_data))

                                        data = ws.receive()
                                        if ws.closed or not check_recv_data(data):
                                            break
                                        opcode, header, data = WSDataPack.unpack(bytes(data))
                                        if opcode != WSOPCode.ACK.value or \
                                            'ack' not in header and \
                                            header['ack'] != WSOPCode.PUSH_CHUNK.value:
                                            break
                                        ws_server.state.update_unack_task_modify_t(file_name)
                                    else:
                                        break
                                    gevent.sleep(0.1)
                                    offset += len(chunk_data)
                            ws_server.cached_fobjects.unref(file_name)
                        except Exception as ex:
                            app.logger.error(ex)
                    else:
                        app.logger.error('chunks is not in header')
            else:
                app.logger.debug('file %s not exits, opcode: %s' % (file_name, str(WSOPCode(opcode))))

            if not send_ack:
                ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': WSOPCode.TASK_NOT_EXIST.value}))

            gevent.sleep(0.1)
    except Exception as ex:
        app.logger.error(ex)

@sockets.route('/websocket/task/pull')
def websocket_task_pull(ws):
    try:
        while not ws.closed:
            data = ws.receive()
            if ws.closed or not check_recv_data(data):
                continue
            opcode, header, _ = WSDataPack.unpack(bytes(data))

            if opcode == WSOPCode.PULL_TASK.value:
                tasks = []
                client_id = ''
                if 'client_id' in header:
                    client_id = header['client_id']
                    while True:
                        try:
                            task = ws_server.state.fetch_task(client_id, sync=False)
                            if task is not None and 'opcode' in task:
                                app.logger.debug('fetch task from client %s, %s %s' % (client_id, str(WSOPCode(task['opcode'])), task['file_name']))
                                tasks.append(task)
                            else:
                                break
                            if len(tasks) > 9:
                                break
                        except Exception as ex:
                            app.logger.error('exception occured while fetch task from %s' % client_id)
                            app.logger.error(ex)
                ws_server.state.sync_tasks()
                ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': opcode, 'tasks': tasks}, None))

            gevent.sleep(0.1)
    except Exception as ex:
        app.logger.error(ex)

@sockets.route('/websocket/task/finish')
def websocket_task_finish(ws):
    try:
        while not ws.closed:
            data = ws.receive()
            if ws.closed or not check_recv_data(data):
                continue
            opcode, header, _ = WSDataPack.unpack(bytes(data))

            if opcode == WSOPCode.FINISH_TASK.value:
                if 'file_name' in header and 'client_id' in header:
                    file_name = header['file_name']
                    client_id = header['client_id']
                    app.logger.debug('task %s to client %s finished' % (file_name, client_id))
                    ws_server.state.finish_task(file_name, client_id)
                    #ws.send(WSDataPack.pack(WSOPCode.ACK.value, {'ack': opcode, 'file_name': file_name, 'client_id': client_id}, None))
                else:
                    app.logger.error('invalid header, file_name or client_id lost')

            gevent.sleep(0.1)
    except Exception as ex:
        app.logger.error(ex)

@sockets.route('/websocket/state/heartbeat')
def websocket_state_heartbeat(ws):
    try:
        while not ws.closed:
            data = ws.receive()
            if ws.closed or not check_recv_data(data):
                continue
            opcode, header, _ = WSDataPack.unpack(bytes(data))
            if 'client_id' in header:
                client_id = header['client_id']
                if ws_server.state.active_client(client_id):
                    ws_server.state.refresh_unack_task(client_id)
    except Exception as ex:
        app.logger.error(ex)

@sockets.route('/websocket/perf/chunk/push')
def websocket_perf_chunk(ws):
    count = 1
    while not ws.closed:
        d = ws.receive()
        print('%s, length: %s' % (count, len(d)))
        count += 1
        ws.send(b'')